from collections.abc import Iterable
from typing import Any, Callable, Generic, Mapping, TypeVar

from zkl_pyutils_fsspec import FsLike, resolve_child_fs

from zkl_aiutils_datasets.basics import Dataset, DatasetWriter
from zkl_aiutils_datasets.loading import load_dataset_by_scripting_from_itself
from zkl_aiutils_datasets.pipelining import RandomDispatchingWriter, RandomGatheringDataset
from zkl_aiutils_datasets.processing import DelegatingDataset
from zkl_aiutils_datasets.scripting import make_dataset_metadata, write_dataset_metadata
from zkl_aiutils_datasets.utils.weight import parse_item_weight_pairs
from .datasets import DirectoryDatasets

AnySample = TypeVar('AnySample', bound=Any)
AnyDataset = TypeVar('AnyDataset', bound=Dataset)
AnyDatasetWriter = TypeVar('AnyDatasetWriter', bound=DatasetWriter)


class RandomGatheringDirectoryDatasets(
    DelegatingDataset[AnySample],
    DirectoryDatasets[AnyDataset],
    Generic[AnySample, AnyDataset],
):
    def __init__(self,
        fs: FsLike,
        child_factory: Callable[[FsLike], AnyDataset] = load_dataset_by_scripting_from_itself,
        random_seed: int = 42,
    ):
        DirectoryDatasets.__init__(self, fs, child_factory)
        DelegatingDataset.__init__(self, RandomGatheringDataset(self.indexed_children, random_seed=random_seed))


class RandomDispatchingDirectoryDatasetsWriter(
    RandomDispatchingWriter[AnySample],
    Generic[AnySample, AnyDatasetWriter],
):
    def __init__(self,
        fs: FsLike,
        names: Mapping[str, float] | Iterable[tuple[str, float]] | Iterable[str],
        child_factory: Callable[[FsLike], AnyDatasetWriter], *,
        loader: Callable | type | None = RandomGatheringDirectoryDatasets,
        random_seed: int = 42,
    ):
        names, weight = parse_item_weight_pairs(names)
        writers = dict((child_factory(resolve_child_fs(fs, name)), weight) for name, weight in zip(names, weight))
        RandomDispatchingWriter.__init__(self, writers, random_seed=random_seed)

        if loader is not None:
            metadata = make_dataset_metadata(loader)
            write_dataset_metadata(fs, metadata)
